
from models import RIAD
from torch.utils.data import DataLoader
from datasets import mvtec
from torch import optim
from utils.schedulers import *
from utils.favae_utils import *

def riad_run(cfg,phase='train',weights=''):
    print('using riad')
    category = cfg['normal_class']
    train_batch_size = cfg['train_batch_size']
    test_batch_size = cfg['test_batch_size']
    load_size = cfg['load_size']
    input_size = cfg['input_size']
    epochs = cfg['epochs']
    scheduler_name = cfg['scheduler']
    test_interval = cfg['test_interval']

    if category == 'all':
        train_class = mvtec.CLASS_NAMES
    else:
        train_class = [category]

    model = RIAD(cfg)

    for c in train_class:
        train_dataset = mvtec.MVTecDataset(
            root_path=cfg['dataset_dir'], class_name=c, is_train=True, resize=load_size, cropsize=input_size)
        train_dataloader = DataLoader(
            train_dataset, batch_size=train_batch_size, pin_memory=False, num_workers=2)
        test_dataset = mvtec.MVTecDataset(
            root_path=cfg['dataset_dir'], class_name=c, is_train=False, resize=load_size, cropsize=input_size)
        test_dataloader = DataLoader(
            test_dataset, batch_size=test_batch_size, pin_memory=False, num_workers=2)

        if phase == 'train':
            for epoch in range(epochs):
                model.train(epoch, train_dataloader ,c)
                if epoch % test_interval == 0:
                    model.test(c, test_dataloader)
                    model.evaluate(c)
                    model.init_results_list()
        elif phase == 'test':
            model.test(c, test_dataloader , weight_path = weights)
            model.evaluate(c)
            model.init_results_list()

